package gummy

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

const (
	wsURL     = "wss://dashscope.aliyuncs.com/api-ws/v1/inference/" // WebSocket服务器地址
	audioFile = "D:/xj/mygin/pkg/gummy/asr_example.wav"             // 替换为您的音频文件路径
)

var dialer = websocket.DefaultDialer

// 定义结构体来表示JSON数据
type Header struct {
	Action       string                 `json:"action"`
	TaskID       string                 `json:"task_id"`
	Streaming    string                 `json:"streaming"`
	Event        string                 `json:"event"`
	ErrorCode    string                 `json:"error_code,omitempty"`
	ErrorMessage string                 `json:"error_message,omitempty"`
	Attributes   map[string]interface{} `json:"attributes"`
}

type Output struct {
	Translations  []Translation `json:"translations"`
	Transcription Transcription `json:"transcription"`
}

type Translation struct {
	SentenceID   int    `json:"sentence_id"`
	BeginTime    int64  `json:"begin_time"`
	EndTime      int64  `json:"end_time"`
	Text         string `json:"text"`
	Lang         string `json:"lang"`
	PreEndFailed bool   `json:"pre_end_failed"`
	Words        []Word `json:"words"`
	SentenceEnd  bool   `json:"sentence_end"`
}

type Transcription struct {
	SentenceID  int    `json:"sentence_id"`
	BeginTime   int64  `json:"begin_time"`
	EndTime     int64  `json:"end_time"`
	Text        string `json:"text"`
	Words       []Word `json:"words"`
	SentenceEnd bool   `json:"sentence_end"`
}

type Word struct {
	BeginTime   int64  `json:"begin_time"`
	EndTime     int64  `json:"end_time"`
	Text        string `json:"text"`
	Punctuation string `json:"punctuation"`
	Fixed       bool   `json:"fixed"`
	SpeakerID   *int   `json:"speaker_id"`
}

type Payload struct {
	TaskGroup  string  `json:"task_group"`
	Task       string  `json:"task"`
	Function   string  `json:"function"`
	Model      string  `json:"model"`
	Parameters Params  `json:"parameters"`
	Input      Input   `json:"input"`
	Output     *Output `json:"output,omitempty"`
}

type Params struct {
	Format                     string   `json:"format"`
	SampleRate                 int      `json:"sample_rate"`
	VocabularyID               string   `json:"vocabulary_id,omitempty"`
	TranslationTargetLanguages []string `json:"translation_target_languages,omitempty"`
	TranscriptionEnabled       bool     `json:"transcription_enabled,omitempty"`
	TranslationEnabled         bool     `json:"translation_enabled,omitempty"`
}

type Input struct {
}

type Event struct {
	Header  Header  `json:"header"`
	Payload Payload `json:"payload"`
}

type TranslationResult struct {
	Transcription string   `json:"transcription"`
	Translations  []string `json:"translations"`
	IsFinal       bool     `json:"is_final"` // 标识是否是最终结果
}

// 连接WebSocket服务
func ConnectWebSocket(apiKey string) (*websocket.Conn, error) {
	header := make(http.Header)
	header.Add("X-DashScope-DataInspection", "enable")
	header.Add("Authorization", fmt.Sprintf("bearer %s", apiKey))
	conn, _, err := dialer.Dial(wsURL, header)
	return conn, err
}

// 启动一个goroutine异步接收WebSocket消息
// 修改后的StartResultReceiver函数
func StartResultReceiver(conn *websocket.Conn, taskStarted chan<- bool, results chan<- TranslationResult, done chan<- struct{}) {
	defer close(done)

	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			log.Println("解析服务器消息失败:", err)
			return
		}

		var event Event
		if err := json.Unmarshal(message, &event); err != nil {
			log.Println("解析事件失败:", err)
			continue
		}

		switch event.Header.Event {
		case "task-started":
			taskStarted <- true

		case "result-generated":
			if event.Payload.Output != nil {
				result := TranslationResult{
					IsFinal: false,
				}

				// 收集识别结果
				if event.Payload.Output.Transcription.Text != "" {
					result.Transcription = event.Payload.Output.Transcription.Text
				}

				// 收集翻译结果
				for _, translation := range event.Payload.Output.Translations {
					result.Translations = append(result.Translations, translation.Text)
				}

				// 发送部分结果
				results <- result
			}

		case "task-finished":
			// 发送最终结果
			results <- TranslationResult{
				IsFinal:       true,
				Transcription: "", // 可以包含最终转录
				Translations:  []string{"任务完成"},
			}
			return

		case "task-failed":
			results <- TranslationResult{
				IsFinal:       true,
				Transcription: "",
				Translations:  []string{fmt.Sprintf("任务失败: %s", event.Header.ErrorMessage)},
			}
			return
		}
	}
}

// 发送run-task指令
func SendRunTaskCmd(conn *websocket.Conn) (string, error) {
	runTaskCmd, taskID, err := generateRunTaskCmd()
	if err != nil {
		return "", err
	}
	err = conn.WriteMessage(websocket.TextMessage, []byte(runTaskCmd))
	return taskID, err
}

// 生成run-task指令
func generateRunTaskCmd() (string, string, error) {
	taskID := uuid.New().String()
	runTaskCmd := Event{
		Header: Header{
			Action:    "run-task",
			TaskID:    taskID,
			Streaming: "duplex",
		},
		Payload: Payload{
			TaskGroup: "audio",
			Task:      "asr",
			Function:  "recognition",
			Model:     "gummy-realtime-v1",
			Parameters: Params{
				Format:                     "wav",
				SampleRate:                 16000,
				TranscriptionEnabled:       true,
				TranslationEnabled:         true,
				TranslationTargetLanguages: []string{"en"},
			},
			Input: Input{},
		},
	}
	runTaskCmdJSON, err := json.Marshal(runTaskCmd)
	return string(runTaskCmdJSON), taskID, err
}

// 等待task-started事件
func WaitForTaskStarted(taskStarted chan bool) {
	select {
	case <-taskStarted:
		fmt.Println("任务开启成功")
	case <-time.After(10 * time.Second):
		log.Fatal("等待task-started超时，任务开启失败")
	}
}

// 发送音频数据
func SendAudioData(conn *websocket.Conn) error {
	file, err := os.Open(audioFile)
	if err != nil {
		return err
	}
	defer file.Close()

	buf := make([]byte, 1024) // 假设100ms的音频数据大约为1024字节
	for {
		n, err := file.Read(buf)
		if n == 0 {
			break
		}
		if err != nil && err != io.EOF {
			return err
		}
		err = conn.WriteMessage(websocket.BinaryMessage, buf[:n])
		if err != nil {
			return err
		}
		time.Sleep(100 * time.Millisecond)
	}
	return nil
}

// 发送finish-task指令
func SendFinishTaskCmd(conn *websocket.Conn, taskID string) error {
	finishTaskCmd, err := generateFinishTaskCmd(taskID)
	if err != nil {
		return err
	}
	err = conn.WriteMessage(websocket.TextMessage, []byte(finishTaskCmd))
	return err
}

// 生成finish-task指令
func generateFinishTaskCmd(taskID string) (string, error) {
	finishTaskCmd := Event{
		Header: Header{
			Action:    "finish-task",
			TaskID:    taskID,
			Streaming: "duplex",
		},
		Payload: Payload{
			Input: Input{},
		},
	}
	finishTaskCmdJSON, err := json.Marshal(finishTaskCmd)
	return string(finishTaskCmdJSON), err
}

// 处理事件
func handleEvent(conn *websocket.Conn, event Event, taskStarted chan<- bool, taskDone chan<- bool) bool {
	switch event.Header.Event {
	case "task-started":
		fmt.Println("收到task-started事件")
		taskStarted <- true
	case "result-generated":
		fmt.Println("服务器返回结果：")
		// 解析 Translations 和 Transcription
		if event.Payload.Output != nil {
			for _, translation := range event.Payload.Output.Translations {
				fmt.Printf("	翻译结果 - Sentence ID：%d, Text：%s\n", translation.SentenceID, translation.Text)
				for _, word := range translation.Words {
					fmt.Printf("	  Word - Begin Time：%d, End Time：%d, Text：%s\n", word.BeginTime, word.EndTime, word.Text)
				}
			}

			transcription := event.Payload.Output.Transcription
			fmt.Printf("	识别结果 - Sentence ID：%d, Text：%s\n", transcription.SentenceID, transcription.Text)
			for _, word := range transcription.Words {
				fmt.Printf("	  Word - Begin Time：%d, End Time：%d, Text：%s\n", word.BeginTime, word.EndTime, word.Text)
			}
		}
	case "task-finished":
		fmt.Println("任务完成")
		taskDone <- true
		return true
	case "task-failed":
		handleTaskFailed(event, conn)
		taskDone <- true
		return true
	default:
		log.Printf("预料之外的事件：%v", event)
	}
	return false
}

// 处理任务失败事件
func handleTaskFailed(event Event, conn *websocket.Conn) {
	if event.Header.ErrorMessage != "" {
		log.Fatalf("任务失败：%s", event.Header.ErrorMessage)
	} else {
		log.Fatal("未知原因导致任务失败")
	}
}

// 关闭连接
func CloseConnection(conn *websocket.Conn) {
	if conn != nil {
		conn.Close()
	}
}
